# encoding:utf-8
"""
@file = jwt_auth
@author = zouju
@create_time = 2023-09-01- 8:20
"""

from rest_framework_simplejwt.exceptions import InvalidToken, AuthenticationFailed
from django.utils.translation import gettext_lazy as _

from system_manage.models import User
from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework_simplejwt.exceptions import AuthenticationFailed
from rest_framework.response import Response


# 自定义的用户模型


class MyJWTAuthentication(JWTAuthentication):
    '''
    自定义JWT认证类，返回自定义User表对象
    自定义get_user()方法即可，授权校验时执行认证类下的authenticate()方法
    '''

    def get_user(self, validated_token):
        try:
            user_id = validated_token['user_id']
        except KeyError:
            raise InvalidToken(_('Token contained no recognizable user identification'))

        try:
            user = User.objects.get(**{'id': user_id})
        except User.DoesNotExist:
            raise AuthenticationFailed(_('User not found'), code='user_not_found')

        return user


def jwt_login_required(function):
    """
    jwt登录认证
    :param function:
    :return:
    """

    def wrapper(request, *args, **kwargs):
        try:
            auth = MyJWTAuthentication()
            user, _ = auth.authenticate(request)
            if user is not None:
                # 验证成功，将用户信息添加到请求对象中
                request.user = user
                return function(request, *args, **kwargs)
            else:
                return Response({'error': 'Invalid token'}, status=401)
        except AuthenticationFailed:
            return Response({'error': 'Invalid token'}, status=401)

    return wrapper
